package tests;

import java.io.FileWriter;
import java.io.IOException;
import java.io.PrintWriter;
import java.util.Arrays;
import java.util.Random;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

import stdp_rules.LogSTDP;
import utils.StateCamera;
import utils.StateCamera.CameraType;
import utils.math_utils.Coordinate;
import Jama.Matrix;

import components.AbstractNeuron;
import components.InputNeuron;
import components.LIFNeuron;
import components.Network;
import components.Synapse;

import connectors.GridConnect;
import executives.RunForwardParallel;

public class ParallelXOR_Test {

	public static final int numXORDist = 15;//Epochs 
	public static final int trainEpochs = 5000;
	public static final int testEpochs = 1000;
	public static final double timeStep = 1.0; //ms
	public static final double epoch = 15; //ms	
	public static final int numInputs = 400;
	public static final int numRes = 1024;
	public static final int memTc = 30;
	private static CountDownLatch latch;
	
	private StateCamera camera;
	private static int [] resDims = {8, 8, 16};
	private static int [] inDims = {20, 20, 1};
	private final Network net;
	private boolean [] trainingIn;
	private boolean [] testingIn;
	private byte [][] teachTargs;
	private byte [][] testTargs;
	private LogSTDP stdp;	
	private LIFNeuron [] res = new LIFNeuron[numRes];
	private InputNeuron [] in = new InputNeuron[numInputs];
	
	public ParallelXOR_Test(){
		net = new Network();
		initialize();
		generateXORPattern();
	}
	
	/**
	 * 
	 */
	private void initialize() {
			
		camera = new StateCamera(CameraType.Bin,
				timeStep, epoch);			
		
		
		in = (Network.initializeInputs(numInputs, 0));
		res =(Network.initializeNeurons(numRes,
				30, 15, 2.5, 0.2, camera));
		

		net.cubicLayout(new Coordinate(0,0,0), inDims ,in);
		net.cubicLayout(new Coordinate(7, 7, 2), resDims,
				res);
		
		GridConnect gCon = new GridConnect();
		Network.initializeSynRandomizers(gCon);
		gCon.stochasticConnect(res, res, 2.25);
		gCon.stochasticConnect(in, res, 16);
		
		stdp = new LogSTDP(5, 10, 0.0003, 0.01f, 3.0f, 30f, 1f, 20.0f);
		
		for(LIFNeuron n : res) {
			n.setStdp(stdp);
			n.setStdpOn(true);
			n.setReporter(false);
		}
		
		for(InputNeuron n : in) {
			n.setStdp(stdp);
		}
		
	}
	
	/**
	 * 
	 * @param teaching
	 */
	public double convertToSpikes(boolean[] inputs,
			double startTime) {
		double clock = startTime;
		double spk;
		double current;
		for(int i = 0, n = inputs.length; i < n; i++) {
			boolean xor = inputs[i];
			
			spk = clock;
			current = clock;
			clock += epoch;
			for(InputNeuron inp : in) {
				
				if(xor && inp.getIndex() % 2 == 0){
					spk = current;
					while(spk < clock) {
						double interval = epoch/5;
						if(inp.peekLast() == null){
							System.out.println("Interval was less than zero");
						}
						spk += interval;
						if(spk < clock) {
							inp.addSpkTime(spk);
						}else{
							break;
						}
					}	
				} 				
			}				
		}
		return clock;
	}	
	
	/**
	 * Generates the XOR test/train inputs/targets
	 */
	public void generateXORPattern() {
		Random rand = new Random();
		trainingIn = new boolean[trainEpochs];
		teachTargs = new byte[trainEpochs - (numXORDist+1)][numXORDist];
		Matrix inXOR = new Matrix(trainEpochs, 1);
		for (int i = 0; i < trainEpochs; i++) {
			trainingIn[i] = rand.nextBoolean();
			inXOR.set(i, 0, trainingIn[i] ? 1 : 0);
			if (i > numXORDist) {
				for (int j = 1; j <= numXORDist; j++) {
					if (i >= j) {
						teachTargs[i-(numXORDist+ 1)][j-1] = (byte)((trainingIn[i] != trainingIn[i-j]) ? 1 : 0);
					}
				}
			}
		}
		
		PrintWriter pw = null;
		try{
			FileWriter fw = new FileWriter("xorIns.dat");
			pw = new PrintWriter(fw);
			inXOR.print(pw, 4, 3);
		} catch (IOException e) {
			e.printStackTrace();
		} finally {
			if(pw != null)
				pw.close();
		}
		
		
		testingIn = new boolean[testEpochs];
		testTargs = new byte[testEpochs - (numXORDist+1)][numXORDist];
		for (int i = 0; i < testEpochs; i++) {
			testingIn[i] = rand.nextBoolean();
			if (i > numXORDist) {
				for (int j = 1; j <= numXORDist; j++) {
					if (i >= j) {
						testTargs[i-(numXORDist+ 1)][j-1] = (byte)((testingIn[i] != testingIn[i-j]) ? 1 : 0);
					}
				}
			}
		}	
		
	}
	
	public Network getNet(){
		return net;
	}
	
	/**
	 * @param args
	 */
	public static void main(String[] args) {
		ParallelXOR_Test test = new ParallelXOR_Test();
		
	
		Matrix wOut;	
		
		try {
			
			//**************************************************************//
			//*************************First Pass***************************//
			System.out.println("First-Pass...");
			test.getStateCamera().readyFilm(trainEpochs, numInputs);
			System.out.println(Runtime.getRuntime().freeMemory()/1000000);
			double endTime = test.convertToSpikes(test.getTrainingIn(), 0);
			System.out.println(Runtime.getRuntime().freeMemory()/1000000);
			System.out.println("Conversion done.");
			latch = new CountDownLatch(1);	
			RunForwardParallel executive = new RunForwardParallel(timeStep, 0,
					endTime, latch, null, test.getAllNeurons());
			ExecutorService pool = Executors.newFixedThreadPool(
					RunForwardParallel.DEFAULT_NUM_THREADS);
			pool.submit(executive).get();
			latch.await();
			test.getStateCamera().printStatesToFile("InputsXOR.dat");
//			ArrayList<Thread> availableThreads =
//					(ArrayList<Thread>)executive.getAvailableThreads();
			System.out.println("Invocations: " +  Synapse.occurances);
			System.out.println("Avg size of queues: " + Synapse.length);
			System.out.println("NaN Activations: " + Synapse.nans);
			//**************************************************************//
			//*************************TRAINING****************************//
			System.out.println("Training...");		
			test.getStateCamera().readyFilm(trainEpochs, numRes);
			for(LIFNeuron n : test.getRes()) {
				n.setStdpOn(false);
				n.setRecorder(test.camera);
				n.setReporter(true);				
			}			
			for(InputNeuron n : test.getIn()) {
				n.setStdpOn(false);
			}			
			double startTime = endTime;
			endTime = test.convertToSpikes(test.getTrainingIn(), startTime);
			latch = new CountDownLatch(1);
			executive = new RunForwardParallel(timeStep, startTime, endTime,
					latch, test.getStateCamera(), test.getAllNeurons());
			
			pool = Executors.newFixedThreadPool(
					RunForwardParallel.DEFAULT_NUM_THREADS);
			pool.submit(executive).get();
			latch.await();
			
			System.out.println("Invocations: " +  Synapse.occurances);
			System.out.println("Avg size of queues: " + Synapse.length);
			System.out.println("NaN Activations: " + Synapse.nans);
			
			//Solving
			System.out.println("Solving...");
//			latch = new CountDownLatch(1);
			double [][] states = Arrays.copyOfRange(test.getStateCamera().
					getCapturedStates(), (numXORDist+ 1),
					test.getStateCamera().getCapturedStates().length);
			Matrix s = new Matrix(states);
			Matrix d = new Matrix(test.getTeachTargs());
			wOut = s.solve(d);
//			latch.countDown();
//			latch.await();
			System.out.println("Network trained.");
			test.getStateCamera().printStatesToFile("train.dat");	
			
			//**************************************************************//
			//*************************TESTING*****************************//
			System.out.println("Testing...");
			test.getStateCamera().readyFilm(testEpochs, numRes);
			startTime = endTime;
			endTime = test.convertToSpikes(test.getTestingIn(), startTime);
			
			latch = new CountDownLatch(1);
			executive = new RunForwardParallel(timeStep, startTime, endTime,
					latch, test.getStateCamera(), test.getAllNeurons());
			pool.submit(executive).get();
			latch.await();
			pool.shutdownNow();

			//Finding Errors
			System.out.println("Error Checking...");
			states = Arrays.copyOfRange(test.getStateCamera().
					getCapturedStates(), (numXORDist+ 1),
					test.getStateCamera().getCapturedStates().length);
			s = new Matrix(states);
			Matrix tOuts = s.times(wOut);
			for(int i = 0; i < tOuts.getRowDimension(); i++) {
				for(int j = 0; j < numXORDist; j++) {
					tOuts.set(i, j, tOuts.get(i,j) >= 0.5 ? 1.0 : 0.0);
				}
			}	
			Matrix dOuts = new Matrix(test.getTestTargs());
			Matrix x = (tOuts.minus(dOuts));
			x.arrayTimesEquals(x);
			
			try {
				FileWriter fw = new FileWriter("TestTargets.dat");
				PrintWriter pw = new PrintWriter(fw);
				dOuts.print(pw, 3, 4);
				pw.close();
				
				FileWriter FW = new FileWriter("TestOuts.dat");
				PrintWriter P = new PrintWriter(FW);
				tOuts.print(P,  3, 4);
				P.close();
				
				
			}catch (IOException e) {
				e.printStackTrace();
			}
			
			double [] errors = new double[numXORDist];
			for (int i = 0; i < numXORDist; i++){
				double val = 0.0;
				for(int j = 0; j < testEpochs - (numXORDist+ 1); j++) {
					val = x.get(j, i) + val;
				}
				errors[i] = val/(double)testEpochs;
			}
			
			System.out.println(Arrays.toString(errors));	
			
			//**************************************************************//
			//*************************RECORDING****************************//		
			test.getStateCamera().printStatesToFile("test.dat");	

			
			
		} catch (Exception e) {
			e.printStackTrace();
		}

	}
	
	public StateCamera getStateCamera(){
		return camera;
	}

	public StateCamera getCamera() {
		return camera;
	}

	public void setCamera(StateCamera camera) {
		this.camera = camera;
	}

	public double [][] getTeachTargs() {
		double [][] tt = new double [teachTargs.length][teachTargs[0].length];
		for(int i = 0; i < teachTargs.length; i++) {
			for(int j = 0; j < teachTargs[i].length; j++) {
				tt[i][j] = (double)teachTargs[i][j];
			}
		}
		return tt;
	}

	public void setTeachTargs(byte[][] teachTargs) {
		this.teachTargs = teachTargs;
	}

	public double [][] getTestTargs() {
		double [][] tt = new double [testTargs.length][testTargs[0].length];
		for(int i = 0; i < testTargs.length; i++) {
			for(int j = 0; j < testTargs[i].length; j++) {
				tt[i][j] = (double)testTargs[i][j];
			}
		}
		return tt;
	}

	public void setTestTargs(byte[][] testTargs) {
		this.testTargs = testTargs;
	}

	public LogSTDP getStdp() {
		return stdp;
	}

	public void setStdp(LogSTDP stdp) {
		this.stdp = stdp;
	}

	public LIFNeuron[] getRes() {
		return res;
	}

	public void setRes(LIFNeuron[] res) {
		this.res = res;
	}

	public InputNeuron[] getIn() {
		return in;
	}

	public void setIn(InputNeuron[] in) {
		this.in = in;
	}

	public boolean[] getTrainingIn() {
		return trainingIn;
	}

	public void setTrainingIn(boolean[] trainingIn) {
		this.trainingIn = trainingIn;
	}

	public boolean[] getTestingIn() {
		return testingIn;
	}

	public void setTestingIn(boolean[] testingIn) {
		this.testingIn = testingIn;
	}

	public AbstractNeuron[] getAllNeurons(){
		AbstractNeuron [] arr = new AbstractNeuron[in.length + res.length];
		for(int i = 0, n = in.length + res.length; i < n; i++) {
			if(i < in.length){
				arr[i] = in[i];
			} else {
				arr[i] = res[i-in.length];
			}
		}
		return arr;
	}
	
}
